import time
import json
import os
import threading
import pyautogui
from pynput import mouse, keyboard
from pynput.mouse import Button
from screenshot import ScreenSelector


recordings_dir = "recordings"
special_key = keyboard.Key.esc  # 使用ESC键作为开始/停止控制键
os.makedirs(recordings_dir, exist_ok=True)


def start_recording(name=None, description=None):
    # 首先让用户选择录屏区域
    print("=" * 60)
    print("【步骤 1/3】选择录制区域")
    print("请选择要录制的屏幕区域（默认全屏）")
    print("提示：双击或按 Enter 确认选择，按 ESC 取消")
    print("=" * 60)
    region_x, region_y, region_w, region_h = ScreenSelector.exec()
    
    events = []
    drag_threshold = 5  # 拖动阈值：移动超过5像素视为拖动

    print("\n" + "=" * 60)
    print("【步骤 2/3】开始录制")
    print(f"请按下 ESC 键开始录制，再次按下 ESC 键结束录制")
    print("录制过程中，您的鼠标点击、拖动和键盘操作将被记录")
    print("=" * 60 + "\n")
    start_event = threading.Event()
    stop_event = threading.Event()
    
    # 拖动状态
    is_dragging = False
    drag_button = None
    drag_start_pos = None
    drag_start_time = None
    drag_points = []
    last_move_time = 0
    move_sample_interval = 0.1  # 移动采样间隔：每100ms记录一个点
    
    # 键盘按下状态跟踪（用于过滤重复事件）
    key_pressed = set()
    record_start_time = None  # 录制开始时间，用于计算相对时间

    def on_key_press(key):
        # ESC键用于控制录制，不记录
        if key == special_key:
            if start_event.is_set():
                print("\n" + "=" * 60)
                print("录制已停止！")
                print("=" * 60 + "\n")
                stop_event.set()
            else:
                print("\n" + "=" * 60)
                print("录制已开始！正在记录您的操作...")
                print("提示：再次按下 ESC 键可结束录制")
                print("=" * 60 + "\n")
                start_event.set()
            return
        
        # 只有在录制开始后才记录键盘事件
        if not start_event.is_set():
            return
        
        # 避免重复记录同一个按键的按下事件
        if key in key_pressed:
            return
        
        # 转换键名为字符串
        try:
            if hasattr(key, 'char') and key.char is not None:
                key_str = key.char
            elif hasattr(key, 'name'):
                key_str = key.name
            else:
                key_str = str(key)
        except:
            key_str = str(key)
        
        key_pressed.add(key)
        timestamp = time.time()
        relative_time = timestamp - record_start_time
        event_index = len(events)
        events.append([event_index, relative_time, "key_press", key_str])
        delay = relative_time - events[-2][1] if len(events) > 1 else 0
        print(f"已记录: 事件#{event_index} | 时间 {relative_time:.3f}s | 延迟 {delay:.3f}s | 按键按下 | {key_str}")
    
    def on_key_release(key):
        # ESC键用于控制录制，不记录
        if key == special_key:
            return
        
        # 只有在录制开始后才记录键盘事件
        if not start_event.is_set():
            return
        
        # 只记录已按下的键的释放事件
        if key not in key_pressed:
            return
        
        key_pressed.discard(key)
        timestamp = time.time()
        
        # 转换键名为字符串
        try:
            if hasattr(key, 'char') and key.char is not None:
                key_str = key.char
            elif hasattr(key, 'name'):
                key_str = key.name
            else:
                key_str = str(key)
        except:
            key_str = str(key)
        
        relative_time = timestamp - record_start_time
        event_index = len(events)
        events.append([event_index, relative_time, "key_release", key_str])
        delay = relative_time - events[-2][1] if len(events) > 1 else 0
        print(f"已记录: 事件#{event_index} | 时间 {relative_time:.3f}s | 延迟 {delay:.3f}s | 按键释放 | {key_str}")

    def on_click(x, y, button, pressed):
        nonlocal is_dragging, drag_button, drag_start_pos, drag_start_time, drag_points, last_move_time
        timestamp = time.time()
        relative_time = timestamp - record_start_time
        
        # 将绝对坐标转换为相对坐标（0-1范围）
        rel_x = (x - region_x) / region_w if region_w > 0 else 0.0
        rel_y = (y - region_y) / region_h if region_h > 0 else 0.0
        # 限制在0-1范围内
        rel_x = max(0.0, min(1.0, rel_x))
        rel_y = max(0.0, min(1.0, rel_y))
        
        if pressed:
            # 鼠标按下
            is_dragging = False
            drag_button = button
            drag_start_pos = (rel_x, rel_y)  # 使用相对坐标
            drag_start_time = timestamp
            drag_points = [(rel_x, rel_y, 0.0)]  # (rel_x, rel_y, relative_time) 相对时间从0开始
            last_move_time = timestamp  # 重置移动时间
            event_index = len(events)
            events.append([event_index, relative_time, "press", button.name, rel_x, rel_y])
            delay = relative_time - events[-2][1] if len(events) > 1 else 0
            print(f"已记录: 事件#{event_index} | 时间 {relative_time:.3f}s | 延迟 {delay:.3f}s | 鼠标按下 | {button.name} | 位置 ({rel_x:.3f}, {rel_y:.3f})")
        else:
            # 鼠标释放
            if drag_start_pos and drag_button == button:
                # 计算相对坐标的距离（以像素为单位）
                abs_start_x = drag_start_pos[0] * region_w + region_x
                abs_start_y = drag_start_pos[1] * region_h + region_y
                distance = ((x - abs_start_x)**2 + (y - abs_start_y)**2)**0.5
                
                # 如果移动距离超过阈值，视为拖动
                if distance > drag_threshold or len(drag_points) > 1:
                    is_dragging = True
                    # 记录释放点的相对时间（相对于按下时刻）
                    drag_relative_time = timestamp - drag_start_time
                    drag_points.append((rel_x, rel_y, drag_relative_time))
                    # 移除之前的press事件，添加拖动事件
                    if events and events[-1][2] == "press" and events[-1][3] == button.name:
                        event_index -= 1  # 回退序号
                        events.pop()
                    event_index = len(events)
                    events.append([event_index, relative_time, "drag", button.name, drag_points])
                    delay = relative_time - events[-2][1] if len(events) > 1 else 0
                    print(f"已记录: 事件#{event_index} | 时间 {relative_time:.3f}s | 延迟 {delay:.3f}s | 鼠标拖动 | {button.name} | {len(drag_points)} 个点 | 持续时间 {drag_relative_time:.3f}s")
                else:
                    # 视为点击
                    if events and events[-1][2] == "press" and events[-1][3] == button.name:
                        # 将press改为click，更新时间戳（保持序号不变）
                        events[-1][1] = relative_time  # 更新相对时间戳
                        events[-1][2] = "click"
                        events[-1][3] = button.name
                        events[-1][4] = rel_x
                        events[-1][5] = rel_y
                        delay = relative_time - events[-2][1] if len(events) > 1 else 0
                        print(f"已记录: 事件#{events[-1][0]} | 时间 {relative_time:.3f}s | 延迟 {delay:.3f}s | 鼠标点击 | {button.name} | 位置 ({rel_x:.3f}, {rel_y:.3f})")
                
                is_dragging = False
                drag_button = None
                drag_start_pos = None
                drag_points = []
                last_move_time = 0

    def on_move(x, y):
        nonlocal is_dragging, drag_points, last_move_time, drag_start_time
        # 仅在拖动时记录移动路径，使用采样间隔避免记录过多点
        if drag_start_pos is not None:
            current_time = time.time()
            if current_time - last_move_time >= move_sample_interval:
                # 将绝对坐标转换为相对坐标（0-1范围）
                rel_x = (x - region_x) / region_w if region_w > 0 else 0.0
                rel_y = (y - region_y) / region_h if region_h > 0 else 0.0
                rel_x = max(0.0, min(1.0, rel_x))
                rel_y = max(0.0, min(1.0, rel_y))
                # 记录坐标和相对时间（相对于按下时刻）
                relative_time = current_time - drag_start_time
                drag_points.append((rel_x, rel_y, relative_time))
                last_move_time = current_time

    with keyboard.Listener(on_press=on_key_press, on_release=on_key_release) as keyboard_listener:
        start_event.wait() 
        record_start_time = time.time()  # 记录录制开始时间
        events.append([0.0, 0, -1])  # 第一个事件时间戳为0（相对时间）
        with mouse.Listener(on_click=on_click, on_move=on_move) as mouse_listener:
            stop_event.wait() 
            mouse_listener.stop()
            keyboard_listener.stop()

    if len(events) > 1:
        print("\n" + "=" * 60)
        print("【步骤 3/3】保存录制")
        print(f"录制完成！共记录了 {len(events) - 1} 个操作事件")
        print("=" * 60 + "\n")
        
        # 如果提供了名称和描述，直接使用；否则通过 input() 获取
        if name:
            script_name = name
        else:
            script_name = input("请输入脚本名称: ").strip()
        
        if not script_name:
            print("录制已取消（未提供脚本名称）。")
            return None, None
        
        if description is None:
            description = input("请输入脚本描述（可选，直接回车跳过）: ").strip()
        
        save_path = os.path.join(recordings_dir, f"{script_name}.json")
        # 直接保存events（跳过第一个开始时间戳），将press转换为click
        data_to_save = {
            "description": description if description else "",
            "events": []
        }
        for event in events[1:]:
            # 如果还有未处理的press事件（录制中断），转换为click
            if event[2] == "press":
                # 保持序号，转换为click
                data_to_save["events"].append([event[0], event[1], "click", event[3], event[4], event[5]])
            else:
                data_to_save["events"].append(event)
        with open(save_path, "w", encoding='utf-8') as f:
            json.dump(data_to_save, f, ensure_ascii=False, indent=2)
        print(f"\n脚本已保存到: {save_path}")
        if description:
            print(f"脚本描述: {description}")
        print("=" * 60 + "\n")
        return script_name, description
    else:
        print("\n录制已取消（未记录任何操作）。\n")
        return None, None


def playback(name, loop_count=1):
    file_path = os.path.join(recordings_dir, f"{name}.json")
    if not os.path.exists(file_path):
        print(f"未找到名为 '{name}' 的脚本。")
        return
    
    with open(file_path, "r", encoding='utf-8') as f: 
        data = json.load(f)
    
    events = data.get("events", [])
    description = data.get("description", "")
    
    if not events:
        print(f"脚本 '{name}' 中没有操作事件。")
        return
    
    # 让用户选择播放区域（可以不同于录制时的区域）
    print("=" * 60)
    print("【播放脚本】选择播放区域")
    print("请选择要播放的屏幕区域（默认全屏）")
    print("提示：双击或按 Enter 确认选择，按 ESC 取消")
    print("=" * 60)
    region_x, region_y, region_w, region_h = ScreenSelector.exec()
        
    print("\n" + "=" * 60)
    print(f"【开始播放】脚本: {name}")
    if description:
        print(f"描述: {description}")
    print(f"播放次数: {loop_count} 次")
    print(f"事件总数: {len(events)} 个")
    print("提示：按 ESC 键可随时停止播放")
    print("=" * 60 + "\n")
    playing = True

    def on_press(key):
        nonlocal playing
        if key == keyboard.Key.esc:
            playing = False
            print("\n播放已暂停。\n")

    # 辅助函数：将字符串键名转换为pynput的Key对象或字符
    def parse_key(key_str):
        """将字符串键名转换为可用的键对象"""
        # 如果是单个字符，直接返回
        if len(key_str) == 1:
            return key_str
        
        # 尝试作为特殊键
        try:
            # 处理像 'ctrl_l', 'shift_l' 这样的键（去掉左右标识）
            if key_str.endswith('_l') or key_str.endswith('_r'):
                base_name = key_str[:-2]
                if hasattr(keyboard.Key, base_name):
                    return getattr(keyboard.Key, base_name)
            
            # 直接尝试键名
            if hasattr(keyboard.Key, key_str):
                return getattr(keyboard.Key, key_str)
            
            # 处理 'Key.xxx' 格式
            if key_str.startswith('Key.'):
                key_name = key_str.replace('Key.', '')
                if hasattr(keyboard.Key, key_name):
                    return getattr(keyboard.Key, key_name)
        except:
            pass
        
        # 如果都找不到，返回原始字符串
        return key_str
        
    # 创建键盘控制器（复用）
    kb_controller = keyboard.Controller()
    
    with keyboard.Listener(on_press=on_press) as keyboard_listener:
        for cnt in range(loop_count):
            if not playing:
                break
            print(f"\n--- 第 {cnt + 1}/{loop_count} 次播放 ---")
            playback_start_time = time.time()  # 记录重放开始时间（绝对时间）
            last_event_time = 0
            
            for event in events:
                if not playing:
                    break
                
                event_index = int(event[0])
                timestamp = event[1]
                event_type = event[2]
                
                # 计算延迟：当前事件时间戳 - 上一个事件时间戳
                delay = timestamp - last_event_time
                if delay > 0:
                    time.sleep(delay)
                
                # 计算绝对时间（从重放开始的时间）
                absolute_time = time.time() - playback_start_time
                last_event_time = timestamp
                
                # 显示序号和绝对时间
                index_str = f"#{event_index}" if event_index is not None else "?"
                
                if event_type == "click":
                    btn = event[3]
                    rel_x = event[4]
                    rel_y = event[5]
                    # 将相对坐标转换为绝对坐标
                    abs_x = region_x + rel_x * region_w
                    abs_y = region_y + rel_y * region_h
                    pyautogui.moveTo(abs_x, abs_y)
                    pyautogui.click(button=btn)
                    print(f"[事件{index_str}] {absolute_time:.3f}s | 点击: {btn} | 位置 ({abs_x:.1f}, {abs_y:.1f}) | 延迟 {delay:.3f}s")

                elif event_type == "drag":
                    btn = event[3]
                    points = event[4]
                    
                    # 记录拖动开始时的绝对时间
                    drag_start_absolute_time = time.time() - playback_start_time
                    
                    # 将相对坐标转换为绝对坐标
                    rel_start_x, rel_start_y, start_t = points[0]
                    start_x = region_x + rel_start_x * region_w
                    start_y = region_y + rel_start_y * region_h
                    pyautogui.moveTo(start_x, start_y)
                    
                    # 记录按下鼠标的实际时间（绝对时间）
                    pyautogui.mouseDown(button=btn)
                    drag_start_real_time = time.time()
                    
                    # 移动操作本身需要时间（约0.1秒），需要提前执行
                    estimated_move_time = 0.1  # 预估每次移动需要的时间
                    
                    for point in points[1:]:
                        if not playing:
                            break
                        
                        rel_x, rel_y, relative_time = point
                        # 将相对坐标转换为绝对坐标
                        abs_x = region_x + rel_x * region_w
                        abs_y = region_y + rel_y * region_h
                        # 计算该点应该在何时到达（绝对时间）
                        target_time = drag_start_real_time + relative_time
                        current_time = time.time()
                        wait_time = target_time - estimated_move_time - current_time
                        
                        # 如果还没到开始移动的时间，等待
                        if wait_time > 0.001:
                            time.sleep(wait_time)
                            move_start_time = time.time()
                        else:
                            move_start_time = current_time

                        pyautogui.moveTo(abs_x, abs_y)

                        current_time = time.time()
                        actual_move_time = current_time - move_start_time
                        estimated_move_time = actual_move_time
                        
                        if (current_time - target_time) > 0.05:
                            print(f"警告: 移动延迟 {current_time - target_time:.3f}s，实际移动时间: {actual_move_time:.3f}s")

                    current_time = time.time()
                    pyautogui.mouseUp(button=btn)
                    rel_end_x, rel_end_y, end_t = points[-1]
                    abs_end_x = region_x + rel_end_x * region_w
                    abs_end_y = region_y + rel_end_y * region_h
                    actual_duration = current_time - drag_start_real_time
                    # 拖动结束时的绝对时间
                    drag_end_absolute_time = time.time() - playback_start_time
                    print(f"[事件{index_str}] {drag_end_absolute_time:.3f}s | 拖动: {btn} | 从 ({start_x:.1f}, {start_y:.1f}) 到 ({abs_end_x:.1f}, {abs_end_y:.1f}) | 延迟 {delay:.3f}s | 预期 {end_t:.3f}s | 实际 {actual_duration:.3f}s")
                
                elif event_type == "key_press":
                    key_str = event[3]
                    key = parse_key(key_str)
                    # 使用pynput模拟按键按下
                    try:
                        kb_controller.press(key)
                        print(f"[事件{index_str}] {absolute_time:.3f}s | 按键按下: {key_str} | 延迟 {delay:.3f}s")
                    except Exception as e:
                        print(f"[事件{index_str}] {absolute_time:.3f}s | 按键按下失败: {key_str} | 错误: {e}")
                
                elif event_type == "key_release":
                    key_str = event[3]
                    key = parse_key(key_str)
                    # 使用pynput模拟按键释放
                    try:
                        kb_controller.release(key)
                        print(f"[事件{index_str}] {absolute_time:.3f}s | 按键释放: {key_str} | 延迟 {delay:.3f}s")
                    except Exception as e:
                        print(f"[事件{index_str}] {absolute_time:.3f}s | 按键释放失败: {key_str} | 错误: {e}")
                
                else:
                    print(f"未知事件格式: {event}")
        keyboard_listener.stop()
        print("\n" + "=" * 60)
        print("播放已结束。")
        print("=" * 60 + "\n")


def main():
    print("=== Mouse Click Recorder & Replayer ===")
    print("Commands:")
    print("- record           : Start a new recording")
    print("- play [name] [n]  : Play saved recording [name] n times")
    print("- exit             : Exit the program")

    while True:
        cmd = input(">> ").strip()
        if cmd == "record":
            start_recording()
        elif cmd.startswith("play"):
            parts = cmd.split()
            if len(parts) == 3 and parts[2].isdigit():
                name = parts[1]
                count = int(parts[2])
                playback(name, count)
            else:
                print("Usage: play [name] [count]")
        elif cmd == "exit":
            break
        else:
            print("Unknown command.")

if __name__ == "__main__":
    main()  